//
// modified by wsx on 2020.12.12
//
#ifndef _YOLO_H_
#define _YOLO_H_

#include "torch/script.h"
#include "torch/torch.h"
#include "ATen/core/TensorBody.h"
#include "opencv2/opencv.hpp"
#include "opencv2/core/mat.hpp"
#include "utils.h"
#include "vector"

using namespace std;

class YOLO  
{
public:

    /**************************************************************************
     * @brief 构造函数
     * @param modelPath 模型路径
     **************************************************************************/ 
    YOLO(const char* modelPath);

    /**************************************************************************
     * @brief 检测图片
     * @param image 待检测的原始图片
     * @return 大小为(n, 6)的向量，代表检测出来的n个目标的位置、置信度、种类信息(ymin, xmin, ymax, xmax, conf, class)
     **************************************************************************/    
    vector<vector<float>> detect_image(cv::Mat image);

    /**************************************************************************
     * @brief 显示检测结果
     * @param output detect_image的输出
     * @param img 要画框的原始图像
     * @param show - 是否显示最终画了框框的图片
     **************************************************************************/   
    void show_results(vector<vector<float>> output, cv::Mat img, bool show=true);

    /**************************************************************************
     * @brief 预处理处理特征层
     * @param input 模型输出的原始特征层
     **************************************************************************/ 
    at::Tensor find_obj(at::Tensor input);

    /**************************************************************************
     * @brief 确定使用的设备类型
     * @return 设备类型
     **************************************************************************/   
    at::DeviceType wether_GPU(void);


private:

    /**************************************************************************
     * @brief 模型变量
     **************************************************************************/   
    torch::jit::script::Module module;

    /**************************************************************************
     * @brief anchors的尺寸，按需修改
     **************************************************************************/   
    float all_anchors[3][3][2] = {{{142, 110}, {192, 243}, {459, 401}},
                                  {{36,   75}, {76,   55}, {72,  146}},
                                  {{12,   16}, {19,   36}, {40,   28}}};

    /**************************************************************************
     * @brief 出入模型的tensor尺寸，按需修改
     **************************************************************************/   
    float model_image_size[2] = {608, 608};

    /**************************************************************************
     * @brief 检测的类别，按需修改
     **************************************************************************/   
    vector<string> class_names = {"urchin", "trepang"};

    /**************************************************************************
     * @brief 设备类型
     **************************************************************************/ 
    at::DeviceType deviceType;
};




#endif

